package main

import (
	"bytes"
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"golang.org/x/tools/go/ast/astutil"
	"io/fs" // Go 1.16+ 用于 WalkDir 和文件权限
	"os"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	// "sync/atomic" // 在当前实现中，Load/Store 更直观，暂时不用 atomic
	// "unsafe" // 配合 atomic 使用，当前不需要
)

// ============================================================================
// 1. 追踪逻辑 (基本无变化)
// ============================================================================

var (
	// goroutineSpace 用于在 runtime.Stack 输出中识别goroutine ID的前缀
	goroutineSpace = []byte("goroutine ")
	// goroutineIndent 使用 sync.Map 存储每个 goroutine 的当前调用栈深度（用于缩进）
	// key: goroutine ID (uint64), value: indent level (int64)
	goroutineIndent sync.Map
)

// 定义ANSI转义序列颜色常量,用于控制台输出着色
const (
	colorReset   = "\033[0m"  // 重置所有格式
	colorGreen   = "\033[32m" // 绿色，用于 Goroutine ID
	colorBlue    = "\033[34m" // 蓝色，用于函数名
	colorMagenta = "\033[35m" // 品红色，用于行号
	// traceFuncName 是我们注入和追踪的函数的名字
	traceFuncName = "Trace" // 硬编码追踪函数名，简化单文件逻辑
)

// curGoroutineID 获取当前 goroutine 的 ID。
// 注意：此方法依赖 runtime.Stack 的输出格式，可能随 Go 版本变化，且有性能开销。
func curGoroutineID() uint64 {
	var buf [64]byte
	n := runtime.Stack(buf[:], false)
	idField := bytes.TrimPrefix(buf[:n], goroutineSpace)
	i := bytes.IndexByte(idField, ' ')
	if i < 0 {
		fmt.Fprintf(os.Stderr, "警告: 获取 goroutine ID 时无法解析堆栈: %q\n", buf[:n])
		return 0
	}
	id, err := strconv.ParseUint(string(idField[:i]), 10, 64)
	if err != nil {
		fmt.Fprintf(os.Stderr, "警告: 解析 goroutine ID '%s' 失败: %v\n", string(idField[:i]), err)
		return 0
	}
	return id
}

// printTrace 打印格式化的跟踪日志（函数进入或退出）
func printTrace(id uint64, name string, line int, arrow string, indent int) {
	indents := strings.Repeat("  ", indent)
	fmt.Printf("%sgid[%05d]%s %s%s %s%s%s:%s%d%s\n",
		colorGreen, id, colorReset,
		indents, arrow,
		colorBlue, name, colorReset,
		colorMagenta, line, colorReset)
}

// Trace 是被注入到目标函数中的追踪函数。
func Trace() func() {
	pc, _, line, ok := runtime.Caller(1)
	if !ok {
		fmt.Fprintf(os.Stderr, "警告: runtime.Caller(1) 失败\n")
		return func() {}
	}
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		fmt.Fprintf(os.Stderr, "警告: runtime.FuncForPC(%v) 失败\n", pc)
		return func() {}
	}
	name := fn.Name()
	gid := curGoroutineID()

	rawIndent, _ := goroutineIndent.LoadOrStore(gid, int64(0))
	indent := rawIndent.(int64)
	printTrace(gid, name, line, "->", int(indent))
	goroutineIndent.Store(gid, indent+1)

	return func() {
		currentRawIndent, loaded := goroutineIndent.Load(gid)
		if !loaded {
			fmt.Fprintf(os.Stderr, "警告: goroutine %d 的缩进信息在退出时丢失 (函数: %s)\n", gid, name)
			printTrace(gid, name, line, "<-", int(indent)) // Fallback to entry indent
			return
		}
		currentIndent := currentRawIndent.(int64)
		printTrace(gid, name, line, "<-", int(indent)) // Use entry indent for alignment
		newExitIndent := currentIndent - 1
		if newExitIndent <= 0 {
			goroutineIndent.Delete(gid)
		} else {
			goroutineIndent.Store(gid, newExitIndent)
		}
	}
}

// ============================================================================
// 2. AST 插桩逻辑 (修复Instrument中的Apply条件)
// ============================================================================

// instrumenter 结构体用于封装 AST 修改逻辑。
type instrumenter struct{}

// NewInstrumenter 创建一个新的 instrumenter 实例。
func NewInstrumenter() *instrumenter {
	return &instrumenter{}
}

// Instrument 方法解析 Go 源文件，并向其中符合条件的函数注入 `defer Trace()()` 语句。
func (a *instrumenter) Instrument(filename string) ([]byte, error) {
	fset := token.NewFileSet()
	curAST, err := parser.ParseFile(fset, filename, nil, parser.ParseComments)
	if err != nil {
		return nil, fmt.Errorf("解析文件 %s 时出错: %w", filename, err)
	}

	var modified bool
	astutil.Apply(curAST, func(c *astutil.Cursor) bool {
		fd, ok := c.Node().(*ast.FuncDecl)
		// **修复点:** 移除 || len(fd.Body.List) == 0 条件。
		// 现在即使函数体为空 (如 func A(){}) 也会尝试添加defer。
		// 跳过非函数声明或没有函数体的情况（如接口方法）。
		if !ok || fd.Body == nil {
			return true
		}
		if a.addDeferStmt(fd) {
			modified = true
		}
		return true
	}, nil)

	if !modified {
		return nil, nil
	}

	var buf bytes.Buffer
	err = format.Node(&buf, fset, curAST)
	if err != nil {
		return nil, fmt.Errorf("格式化 %s 的新代码时出错: %w", filename, err)
	}
	return buf.Bytes(), nil
}

// addDeferStmt 尝试向函数声明 (fd) 的函数体开头添加 "defer Trace()()" 语句。
func (a *instrumenter) addDeferStmt(fd *ast.FuncDecl) (added bool) {
	// 函数体必须存在才能添加语句
	if fd.Body == nil {
		return false
	}
	// 如果函数体原本是空的 (List为nil)，需要初始化一下
	if fd.Body.List == nil {
		fd.Body.List = make([]ast.Stmt, 0, 1) // 初始化为空切片
	}
	stmts := fd.Body.List

	// 检查是否已存在 "defer Trace()()"
	for _, stmt := range stmts {
		ds, ok := stmt.(*ast.DeferStmt)
		if !ok {
			continue
		}
		outerCall, ok := ds.Call.Fun.(*ast.CallExpr)
		if !ok {
			continue
		}
		ident, ok := outerCall.Fun.(*ast.Ident)
		if !ok {
			continue
		}
		if ident.Name == traceFuncName {
			return false // 已存在
		}
	}

	// 不存在，构造新的 defer 语句
	newDeferStmt := &ast.DeferStmt{
		Call: &ast.CallExpr{
			Fun: &ast.CallExpr{
				Fun: &ast.Ident{Name: traceFuncName},
			},
		},
	}

	// 插入到列表开头
	newList := make([]ast.Stmt, 0, len(stmts)+1)
	newList = append(newList, newDeferStmt)
	newList = append(newList, stmts...)
	fd.Body.List = newList

	return true // 成功添加
}

// ============================================================================
// 3. 命令行处理逻辑 (基本无变化)
// ============================================================================

var (
	// write 标志决定是否将修改写回源文件
	write = flag.Bool("w", false, "将结果写入源文件而不是标准输出")
)

// usage 打印命令行用法信息
func usage() {
	fmt.Fprintf(os.Stderr, "用法: go run instrument.go [-w] [path]\n")
	fmt.Fprintf(os.Stderr, "       instrument [-w] [path] (编译后)\n\n")
	fmt.Fprintf(os.Stderr, "  path: 要处理的 Go 文件或目录。\n")
	fmt.Fprintf(os.Stderr, "        默认当前目录。递归处理 .go 文件 (跳过 _test.go, vendor/, .*)。\n")
	fmt.Fprintf(os.Stderr, "选项:\n")
	flag.PrintDefaults()
}

// main 函数是程序的入口点
func main() {
	flag.Usage = usage
	flag.Parse()

	args := flag.Args()
	targetPath := "."
	if len(args) > 0 {
		targetPath = args[0]
	}

	absTargetPath, err := filepath.Abs(targetPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "错误: 获取绝对路径 '%s' 失败: %v\n", targetPath, err)
		os.Exit(1)
	}

	targetInfo, err := os.Stat(absTargetPath)
	if err != nil {
		if os.IsNotExist(err) {
			fmt.Fprintf(os.Stderr, "错误: 路径 '%s' 不存在\n", absTargetPath)
		} else {
			fmt.Fprintf(os.Stderr, "错误: 访问路径 '%s' 时出错: %v\n", absTargetPath, err)
		}
		os.Exit(1)
	}

	ins := NewInstrumenter()

	// --- 获取工具自身路径 ---
	selfPath := ""
	if execPath, execErr := os.Executable(); execErr == nil {
		selfPath, _ = filepath.Abs(execPath)
	}
	runFilePath := ""
	_, runFile, _, runOk := runtime.Caller(0)
	if runOk {
		runFilePath, _ = filepath.Abs(runFile)
	}
	// --- 工具路径获取结束 ---

	stats := struct{ processed, modified, skipped, errors int }{}

	if targetInfo.IsDir() {
		fmt.Printf("递归处理目录: %s\n", absTargetPath)
		err := filepath.WalkDir(absTargetPath, func(path string, d fs.DirEntry, walkErr error) error {
			if walkErr != nil {
				fmt.Fprintf(os.Stderr, "警告: 访问路径 '%s' 出错: %v, 跳过\n", path, walkErr)
				stats.errors++
				return nil
			}
			absPath, pathErr := filepath.Abs(path)
			if pathErr != nil {
				fmt.Fprintf(os.Stderr, "警告: 获取绝对路径 '%s' 失败: %v, 跳过\n", path, pathErr)
				stats.errors++
				return nil
			}

			if d.IsDir() {
				dirName := d.Name()
				if dirName == "vendor" || (dirName != "." && dirName != ".." && strings.HasPrefix(dirName, ".")) {
					fmt.Printf("跳过目录: %s\n", path)
					stats.skipped++
					return filepath.SkipDir
				}
				return nil
			}
			if !d.Type().IsRegular() {
				stats.skipped++
				return nil
			}
			fileName := d.Name()
			if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
				if strings.HasSuffix(fileName, ".go") {
					stats.skipped++
				}
				return nil
			}
			if absPath == selfPath || absPath == runFilePath {
				fmt.Printf("跳过工具自身文件: %s\n", path)
				stats.skipped++
				return nil
			}

			fmt.Printf("处理文件: %s\n", path)
			stats.processed++
			modified, procErr := processFile(path, ins, *write)
			if procErr != nil {
				fmt.Fprintf(os.Stderr, "  错误: 处理 %s 失败: %v\n", path, procErr)
				stats.errors++
			} else if modified {
				writeMsg := ""
				if *write {
					writeMsg = "并写入"
				}
				fmt.Printf("  成功: 文件 %s 已插桩%s。\n", path, writeMsg)
				stats.modified++
			} else {
				fmt.Printf("  信息: 文件 %s 无需修改。\n", path)
			}
			return nil
		})
		if err != nil {
			fmt.Fprintf(os.Stderr, "错误: 遍历目录 '%s' 时发生意外错误: %v\n", absTargetPath, err)
			stats.errors++
		}
	} else {
		fileName := targetInfo.Name()
		absPath := absTargetPath
		if !strings.HasSuffix(fileName, ".go") || strings.HasSuffix(fileName, "_test.go") {
			if strings.HasSuffix(fileName, "_test.go") {
				fmt.Printf("跳过测试文件: %s\n", absPath)
				stats.skipped++
			} else {
				fmt.Fprintf(os.Stderr, "错误: '%s' 不是有效的 Go 源文件。\n", absPath)
				stats.errors++
			}
		} else if absPath == selfPath || absPath == runFilePath {
			fmt.Printf("跳过工具自身文件: %s\n", absPath)
			stats.skipped++
		} else {
			fmt.Printf("处理单个文件: %s\n", absPath)
			stats.processed++
			modified, procErr := processFile(absPath, ins, *write)
			if procErr != nil {
				fmt.Fprintf(os.Stderr, "  错误: 处理 %s 失败: %v\n", absPath, procErr)
				stats.errors++
			} else if modified {
				writeMsg := ""
				if *write {
					writeMsg = "并写入"
				}
				fmt.Printf("  成功: 文件 %s 已插桩%s。\n", absPath, writeMsg)
				stats.modified++
			} else {
				fmt.Printf("  信息: 文件 %s 无需修改。\n", absPath)
			}
		}
	}

	fmt.Println("\n--- 处理完成 ---")
	fmt.Printf("已处理 Go 文件数: %d\n", stats.processed)
	fmt.Printf("成功修改文件数:  %d\n", stats.modified)
	fmt.Printf("跳过文件/目录数: %d\n", stats.skipped)
	fmt.Printf("处理失败文件数:  %d\n", stats.errors)

	if stats.errors > 0 {
		os.Exit(1)
	}
}

// processFile 函数负责处理单个文件的插桩逻辑。
func processFile(path string, ins *instrumenter, writeToFile bool) (bool, error) {
	newSrc, err := ins.Instrument(path)
	if err != nil {
		return false, fmt.Errorf("插桩分析失败: %w", err)
	}
	if newSrc == nil {
		return false, nil
	}

	if !writeToFile {
		fmt.Println("  (检测到修改，但未写入文件。使用 -w 选项写入)")
	} else {
		originalInfo, statErr := os.Stat(path)
		perms := fs.FileMode(0644) // 默认权限
		if statErr == nil {
			perms = originalInfo.Mode().Perm()
		} else {
			fmt.Fprintf(os.Stderr, "警告: 无法获取文件 %s 的原始权限: %v, 将使用默认权限 0644\n", path, statErr)
		}
		err = os.WriteFile(path, newSrc, perms)
		if err != nil {
			return true, fmt.Errorf("写入文件 '%s' 失败: %w", path, err)
		}
	}
	return true, nil
}
